"""
/* Copyright 2018 The Enflame Tech Company. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
"""
# !/usr/bin/python
# coding=utf-8

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

from tensorflow.python.client import device_lib
from collections import OrderedDict
import subprocess
import sys
sys.path.append("../")
from utils.dtu_logger import LOGGER as logger


class HWInfo(object):
    def __init__(self, device):
        self.device = device
        self.available_device_list = self.get_available_device()
        self.available_device_num = self.get_available_device_num()
        self.num_cluster = self.get_cluster_num()
        self.num_cards = self.get_cards_num()
        if self.device.upper() is 'XLA_DTU':
            assert self.available_device_num == self.num_cards \
                   or self.available_device_num == self.num_cards * 5, \
                'Devices num and cards num not match for DTU'
        self.dict = OrderedDict([('num_cluster', self.num_cluster),
                                 ('num_cards', self.num_cards)])

    def get_available_device(self):
        local_device_protos = device_lib.list_local_devices()
        if self.device.upper() == 'DTU':
            self.device = 'XLA_DTU'
        device_list = [x.name for x in local_device_protos if x.device_type == self.device.upper()]
        if device_list == []:
            raise ValueError("No available {} device".format(self.device))
        else:
            return device_list

    def get_available_device_num(self):
        device_list = self.get_available_device()
        available_device_num = len(device_list)
        logger.info("Detected available device number: {}".format(available_device_num))
        return available_device_num

    def get_cluster_num(self):
        if self.device.lower() == 'xla_dtu':
            try:
                proc = subprocess.Popen(["cat /sys/module/dtu/parameters/ip_mask"],
                                        stdout=subprocess.PIPE, shell=True)
                ip_mask = proc.communicate()[0].strip()
                num_cluster = 0
                num_sip = 0
                rc = proc.returncode
                if rc == 0:
                    ip_mask = int(ip_mask)
                    logger.info("ip=%#x" % ip_mask)
                    for bit_val in range(15, 19):
                        bit_mask = ip_mask >> bit_val & 1
                        if bit_mask == 1:
                            num_cluster += 1
                    for bit_sip in range(34, 42):
                        bit_mask = ip_mask >> bit_sip & 1
                        if bit_mask == 1:
                            num_sip += 1
                    logger.info(
                        "num_cluster:{}, num_sip:{}, total:{}".format(num_cluster, num_sip, num_cluster * num_sip))
                    self.num_sip = num_sip
                    return num_cluster
                else:
                    raise FileNotFoundError("/sys/module/dtu/parameters/ip_mask not Found")
            except Exception as ex:
                logger.warn("Failed to detect ip layout on {}".format(self.device.lower()))
                return 1
        else:
            return 1

    def get_dtu_cards_num(self):
        try:
            proc = subprocess.Popen(["ls /dev |grep dtu |wc -l"], stdout=subprocess.PIPE, shell=True)
            num_dtu_cards = proc.communicate()[0].strip()
            rc = proc.returncode
            if rc == 0:
                logger.info("find {} dtu cards".format(num_dtu_cards))
                return int(num_dtu_cards)
        except Exception as ex:
            logger.warn("Failed to get number of dtu cards.")
            return 0

    def get_gpu_cards_num(self):
        try:
            proc = subprocess.Popen(["nvidia-smi -L |wc -l"], stdout=subprocess.PIPE, shell=True)
            num_gpu_cards = proc.communicate()[0].strip()
            rc = proc.returncode
            if rc == 0:
                return int(num_gpu_cards)
        except Exception as ex:
            logger.warn("Failed to get number of gpu cards.")
            return 0

    def get_cards_num(self):
        if self.device.lower() == 'xla_dtu':
            cards_num = self.get_dtu_cards_num()
            logger.info("Detected {} dtu cards".format(cards_num))
            return cards_num
        elif self.device.lower() in ['gpu', 'xla_gpu']:
            cards_num = self.get_gpu_cards_num()
            logger.info("Detected {} gpu cards".format(cards_num))
            return cards_num
        else:
            logger.warn("No accelerator cards found or selected")
            return 1


if __name__ == '__main__':
    # hw_info = HWInfo('cpu')
    hw_info = HWInfo('dtu')
    print(hw_info.get_available_device())
    print(hw_info.dict)
